package main

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"strconv"
	"strings"
)

// GpuInfo 结构体用于存储 GPU 监控数据
type GpuInfo struct {
	Index       string
	Name        string
	MemoryTotal uint64
	MemoryUsed  uint64
	MemoryFree  uint64
	Temperature float64
	BusId       string
}

func OSExecCmd(cmd string) (string, error) {
	c := exec.Command("bash", "-c", cmd)
	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		return "", err
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		return "", err
	}
	if err := c.Start(); err != nil {
		return "", err
	}
	stdout, err := ioutil.ReadAll(stdoutPipe)
	if err != nil {
		return "", err
	}
	stderr, err := ioutil.ReadAll(stderrPipe)
	if err != nil {
		return "", err
	}
	if err := c.Wait(); err != nil {
		msg := fmt.Sprint(err) + ": " + string(stderr)
		return "", fmt.Errorf(msg)
	}
	return strings.TrimRight(string(stdout), "\n"), nil
}
func ExecCmd(cmd string) (string, error) {
	stdout, err := OSExecCmd(cmd)
	if err != nil {
		return "", err
	}
	return strings.TrimRight(stdout, "\n"), nil
}

func GetGpuInfo() ([]GpuInfo, error) {
	_, err := exec.LookPath("nvidia-smi")
	if err != nil {
		return nil, fmt.Errorf("nvidia-smi not found")
	}
	output, err := ExecCmd("nvidia-smi --query-gpu=index,name,memory.total,memory.used,memory.free,temperature.gpu,pci.bus_id --format=csv,noheader,nounits")
	if err != nil {
		return nil, err
	}
	lines := strings.Split(output, "\n")
	var gpuInfoList []GpuInfo
	for _, line := range lines {
		if line == "" {
			continue
		}
		fields := strings.Split(line, ",")
		if len(fields) != 7 {
			return nil, fmt.Errorf("unexpected output format from nvidia-smi")
		}
		memoryTotal, err := strconv.ParseUint(strings.TrimSpace(fields[2]), 10, 64)
		if err != nil {
			return nil, err
		}
		memoryUsed, err := strconv.ParseUint(strings.TrimSpace(fields[3]), 10, 64)
		if err != nil {
			return nil, err
		}
		memoryFree, err := strconv.ParseUint(strings.TrimSpace(fields[4]), 10, 64)
		if err != nil {
			return nil, err
		}
		temperature, err := strconv.ParseFloat(strings.TrimSpace(fields[5]), 64)
		if err != nil {
			return nil, err
		}

		gpuInfo := GpuInfo{
			Index:       strings.TrimSpace(fields[0]),
			Name:        strings.TrimSpace(fields[1]),
			MemoryTotal: memoryTotal,
			MemoryUsed:  memoryUsed,
			MemoryFree:  memoryFree,
			Temperature: temperature,
			BusId:       fields[6],
		}
		gpuInfoList = append(gpuInfoList, gpuInfo)
	}

	return gpuInfoList, nil
}

func main1() {
	gpuInfoList, err := GetGpuInfo()
	if err != nil {
		fmt.Println(err)
	}
	for _, gpuInfo := range gpuInfoList {
		GpuMemoryTotalInBytes := float64(gpuInfo.MemoryTotal) * 1024 * 1024
		GpuMemoryUsedInBytes := float64(gpuInfo.MemoryUsed) * 1024 * 1024
		GpuMemoryFreeInBytes := float64(gpuInfo.MemoryFree) * 1024 * 1024
		GpuMemoryUsage := float64(gpuInfo.MemoryUsed) / float64(gpuInfo.MemoryTotal) * 100
		GpuTemperature := gpuInfo.Temperature
		println(gpuInfo.Index, gpuInfo.Name, GpuMemoryTotalInBytes, GpuMemoryUsedInBytes, GpuMemoryFreeInBytes, GpuMemoryUsage, GpuTemperature)
	}

}
